import torch
import random
import numpy as np
from hmmlearn.hmm import MultinomialHMM
from sklearn.metrics import classification_report


def load_dict(dict_path):
    vocab = {}
    i = 0
    for line in open(dict_path, 'r', encoding='utf-8'):
        key = line.strip('\n')
        vocab[key] = i
        i += 1
    return vocab, {v: k for k, v in vocab.items()}


train_data = torch.load('deep_learning/train_data')
test_data = torch.load('deep_learning/test_data')
word2id, id2word = load_dict('deep_learning/word2id.txt')
tag2id, id2tag = load_dict('deep_learning/tag2id.txt')


def cal_init_array(total_hidden_seq, hidden_dict):
    print(hidden_dict)
    pi = [0 for i in hidden_dict]
    for _ in total_hidden_seq:
        if _:
            pi[_[0]]+=1
    stat_array = pi
    prob_array = [i/sum(pi) for i in pi]
    return prob_array


def cal_transfer_matrix(total_hidden_seq, hidden_dict):
    matrix = np.zeros([len(hidden_dict),len(hidden_dict)])
    for i in total_hidden_seq:
        if i:
            for j in range(len(i)-1):
                matrix[i[j],i[j+1]]+=1
    for i in range(len(matrix)):
        matrix[i] = np.divide(matrix[i], sum(matrix[i]))
    return matrix

def cal_emmit_matrix(total_obser_seq, total_hidden_seq, hidden_dict, word_dict):
    matrix = np.zeros([len(hidden_dict), len(word_dict)])
    for hid_seq, obs_seq in zip(total_hidden_seq, total_obser_seq):
        for hid, obs in zip(hid_seq, obs_seq):
            matrix[hid][obs]+=1
    for i in range(len(matrix)):
        matrix[i] = np.divide(matrix[i], sum(matrix[i]))
    return matrix

trainhid = [i[1] for i in train_data]
trainobs = [i[0] for i in train_data]


seg_hmm = MultinomialHMM(n_components=5)
seg_hmm.startprob_ = cal_init_array(trainhid, tag2id)
seg_hmm.transmat_ = cal_transfer_matrix(trainhid, tag2id)
seg_hmm.emissionprob_ = cal_emmit_matrix(trainobs, trainhid, tag2id, word2id)

def hmm_decode(hmm_model, word_list, word_dict, hidden_dict):
    
    input = np.reshape(np.array(word_list),(len(np.array(word_list)),1))
    result = hmm_model.decode(input, algorithm='viterbi')   
    sentence = [word_dict[i] for i in word_list]
    segment = [hidden_dict[i] for i in result[1]]
    return sentence, segment


input = '中国建筑市场近年来对外开放步伐进一步加快。'
sen, seg = hmm_decode(seg_hmm, [word2id[i] for i in input], id2word, id2tag)
print(sen, seg)